use std::io::Write;
use std::fs::{self, File};
use std::collections::HashMap;




// --------------------- Syntax --------------------------
#[derive(Debug)]
enum Command {
    A {sym: String},     // @xxx
    L {sym: String},     // (xxx)
    C {dest: String, comp: String, jump: String}, // M=M+D // 0;jmp
}



// ---------------------- Parser ---------------------------
struct Parser {
    tokstream: Vec<String>,
}

impl Parser {

    fn new(file: &str) -> Self {
        let content = fs::read_to_string(file).unwrap();
        let tokstream = content.split("\r\n")
            .filter(|x| !(x.starts_with("//") || x.trim() == ""))
            .map(|x| x.trim().to_string()).collect();
        Parser { tokstream }
    }

    fn cleanup(line: String) -> String {
        // every line may contains extra space and comments
        let mut newline = String::new();
        for i in 0..line.len() {
            let strline = line.as_bytes();
            let space = ' ' as u8;
            let slash = '/' as u8;
            if strline[i] == space {
                continue;
            } else if strline[i] == slash && strline[i+1] == slash {
                break;
            } else {
                newline.push(strline[i] as char);
            }
        }
        return newline;
    }

    pub fn parse(self) -> Vec<Command> {
        let mut cmds = vec![];
        for mut line in self.tokstream {
            line = Self::cleanup(line);
            if line.as_str().starts_with("@") {
                let a = Command::A {sym: line.split_off(1)};
                cmds.push(a);
            } else if line.as_str().starts_with("(") {
                line.pop(); // remove ')'
                let l = Command::L {sym: line.split_off(1)};
                cmds.push(l);
            } else { // must be a compute cmd
                if line.contains(";") {
                    let parts: Vec<&str> = line.split(";").collect();
                    let c = Command::C {
                        dest: String::from("null"),
                        comp: parts[0].to_string(),
                        jump: parts[1].to_string(),
                    };
                    cmds.push(c);
                } else if line.contains("=") {
                    let parts: Vec<&str> = line.split("=").collect();
                    let c = Command::C {
                        dest: parts[0].to_string(),
                        comp: parts[1].to_string(),
                        jump: String::from("null"),
                    };
                    cmds.push(c);
                } else {
                    panic!("Bad C-Command Format!");
                }
            }
        }
        return cmds;
    }
}


// ---------------------- CodeMap ---------------------------
macro_rules! hashmap {
    ( $( $key:expr => $val:expr ),* ) => {
        {
            let mut map = std::collections::HashMap::new();
            $( map.insert( $key, $val ); )*
            map
        }
    };
}

struct CodeMap<'a> {
    dest: HashMap<&'a str, &'a str>,
    comp: HashMap<&'a str, &'a str>,
    jump: HashMap<&'a str, &'a str>,
}

impl<'a> CodeMap<'a> {
    fn new() -> Self {
        let dest = hashmap! (
            "null" => "000", "M"  => "001", "D"  => "010", "MD"  => "011",
            "A"    => "100", "AM" => "101", "AD" => "110", "AMD" => "111"
        );
        let comp = hashmap! {
            "0"   => "0101010", 
            "1"   => "0111111",
            "-1"  => "0111010",
            "D"   => "0001100",
            "A"   => "0110000",   "M"   => "1110000",
            "!D"  => "0001101",
            "!A"  => "0110001",   "!M"  => "1110001",
            "-D"  => "0001111",
            "-A"  => "0110011",   "-M"  => "1110011",
            "D+1" => "0011111",
            "A+1" => "0110111",   "M+1" => "1110111",
            "D-1" => "0001110",
            "A-1" => "0110010",   "M-1" => "1110010",
            "D+A" => "0000010",   "D+M" => "1000010",
            "D-A" => "0010011",   "D-M" => "1010011",
            "A-D" => "0000111",   "M-D" => "1000111",
            "D&A" => "0000000",   "D&M" => "1000000",
            "D|A" => "0010101",   "D|M" => "1010101"
        };
        let jump = hashmap! {
            "null" => "000", "JGT" => "001", "JEQ" => "010", "JGE" => "011",
            "JLT"  => "100", "JNE" => "101", "JLE" => "110", "JMP" => "111"
        };
        CodeMap { dest, comp, jump }
    }

    fn lookup_dest(&self, dest: &str) -> &str {
        match self.dest.get(dest) {
            Some(&bin) => bin,
            None => panic!("No Such a dest: `{}`", dest),
        }
    }

    fn lookup_comp(&self, comp: &str) -> &str {
        match self.comp.get(comp) {
            Some(&bin) => bin,
            None => panic!("No Such a comp: `{}`", comp),
        }
    }

    fn lookup_jump(&self, jump: &str) -> &str {
        match self.jump.get(jump) {
            Some(&bin) => bin,
            None => panic!("No Such a jump: `{}`", jump),
        }
    }

    fn form_address(&self, addr: u16) -> String {
        let bin = format!("{:b}", addr);
        let mut prefix = "0000000000000000"[..16-bin.len()].to_string();
        prefix.push_str(&bin);
        assert!(prefix.len() == 16);
        return prefix;
    }
}


struct SymTable {
    table: HashMap<String, u16>,
    pointer: u16,
}

impl SymTable {
    fn new() -> Self {
        let predefined = hashmap! (
            String::from("SP")   => 0x0000, String::from("LCL")  => 0x0001, String::from("ARG") => 0x0002,
            String::from("THIS") => 0x0003, String::from("THAT") => 0x0004,
            String::from("R0")   => 0x0000, String::from("R1")   => 0x0001, String::from("R2")  => 0x0002,
            String::from("R3")   => 0x0003, String::from("R4")   => 0x0004, String::from("R5")  => 0x0005,
            String::from("R6")   => 0x0006, String::from("R7")   => 0x0007, String::from("R8")  => 0x0008,
            String::from("R9")   => 0x0009, String::from("R10")  => 0x000a, String::from("R11") => 0x000b,
            String::from("R12")  => 0x000c, String::from("R13")  => 0x000d, String::from("R14") => 0x000e,
            String::from("R15")  => 0x000f,
            String::from("SCREEN") => 0x4000, String::from("KBD") => 0x6000
        );
        let pointer = 0x0010;
        SymTable {table: predefined, pointer}
    }

    fn add_variable(&mut self, sym: String) {
        self.table.insert(sym, self.pointer);
        self.pointer += 1;
    }

    fn add_label(&mut self, sym: String, addr: u16) {
        self.table.insert(sym, addr);
    }

    fn contains(&self, sym: &str) -> bool {
        self.table.contains_key(sym)
    }

    fn lookup(&self, sym: &str) -> u16 {
        match self.table.get(sym) {
            Some(&addr) => addr,
            None => panic!("No such a symbol: `{}`", sym),
        }
    }
}

fn is_integer(sym: &str) -> bool {
    for c in sym.chars() {
        if !c.is_digit(10) { return false; }
    }
    return true;
}

fn assembler(path: &str) -> std::io::Result<()> {
    let parser = Parser::new(path);
    let cmds = parser.parse();
    let codemap = CodeMap::new();
    // first pass, create symtable, note that we remove the L cmd
    let mut symtable = SymTable::new();
    let mut newcmds = vec![];
    let mut rom_addr = 0;
    for cmd in cmds.into_iter() {
        match cmd {
            Command::L {sym} => symtable.add_label(sym.to_string(), rom_addr),
            e => {
                newcmds.push(e);
                rom_addr += 1;
            }
        }
    }
    // second pass, resolve symbol
    let bincode = newcmds.into_iter().map(|cmd| 
        match cmd {
            Command::A {sym} => {
                // integer address
                if is_integer(&sym) {
                    let addr = sym.parse().unwrap();
                    return codemap.form_address(addr);
                }
                // labels
                if symtable.contains(&sym) {
                    let addr = symtable.lookup(&sym); 
                    return codemap.form_address(addr);
                }
                // variables
                let addr = symtable.pointer;
                symtable.add_variable(sym);
                return codemap.form_address(addr);
            } 
            Command::C {dest, comp, jump} => {
                let bin_dest = codemap.lookup_dest(&dest);
                let bin_comp = codemap.lookup_comp(&comp);
                let bin_jump = codemap.lookup_jump(&jump);
                format!("111{}{}{}", bin_comp, bin_dest, bin_jump)
            }
            _ => unreachable!(),
        }
    ).collect();
    let hack_path = path.replace(".asm", ".hack");
    output(&hack_path, bincode)?;
    Ok(())
}

fn output(filename: &str, bincode: Vec<String>) -> std::io::Result<()> {
    let mut file = File::create(filename)?;
    for code in bincode {
        file.write(code.as_bytes())?;
        file.write(b"\n")?;
    }
    Ok(())
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 2 {
        println!("Usage: ./mas /path/to/some.asm");
    } else {
        assembler(&args[1]);
    }
}
